/*
    Copyright 2007-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file 

    Declaration of the mcrx_stage template and the concrete stage
    classes. The concrete classes are defined in the files
    xxx_stage.cc and the mcrx_stage helper functions in
    mcrx-stage-impl.h \ingroup mcrx */

// $Id$

#ifndef __mcrx_stage__
#define __mcrx_stage__

#include "mcrx.h"
#include <string>
#include "CCfits/CCfits"
#include "aux_grid.h"
#include "emission-fits.h"
#include "full_sed_grid.h"
#include "ir_grid.h"
#include "aux_emergence.h"
#include "full_sed_emergence.h"
#include "boost/shared_ptr.hpp"
#include "dust_model.h"
#include "scatterer.h"
#include "grain_model.h"
#include "shoot.h"

namespace mcrx {
  template <typename> class dummy_grid_template {
  public:typedef int T_data;};

  // stage class general declaration, followed by concrete stage types
  template <template <typename> class grid_type, typename stage_type> class mcrx_stage;
  template <template <typename> class grid_type> class aux_grid_stage;
  class aux_particle_stage;
  class nonscatter_stage;
  template <template <typename> class grid_type> class scatter_stage;
  template <template <typename> class grid_type> class intensity_stage;
  template <template <typename> class grid_type> class ir_stage;
  template <template <typename> class grid_type> class ir_intensity_stage;
  template <template <typename> class grid_type> class thick_ir_stage;

  // stage types general declaration, followed by specializations for all our types
  template <template<typename> class, typename> class stage_types;

  template <template <typename> class grid_type> 
  class stage_types<grid_type, aux_grid_stage<grid_type> >;
  
  template <> class stage_types<dummy_grid_template, aux_particle_stage>;

  template <> class stage_types<dummy_grid_template, nonscatter_stage>;

  template <template <typename> class grid_type> 
  class stage_types<grid_type, scatter_stage<grid_type> >;

  template <template <typename> class grid_type> 
  class stage_types<grid_type, intensity_stage<grid_type> >;

  template <template <typename> class grid_type> 
  class stage_types<grid_type, ir_stage<grid_type> >;

  template <template <typename> class grid_type> 
  class stage_types<grid_type, ir_intensity_stage<grid_type> >;

  template <template <typename> class grid_type> 
  class stage_types<grid_type, thick_ir_stage<grid_type> >;

  template<typename T_scatterer_vector>
  T_scatterer_vector
  read_dust_grains(const Preferences& p, const T_unit_map& units);

  class Mcrx;
  void execute_adaptive_stages(Mcrx&, const Preferences& p);
  void execute_arepo_stages(Mcrx&, const Preferences& p);
};

/** This is a traits class that contains the types needed for the
    different stages. */
template <template<typename> class, typename> class mcrx::stage_types {};


/** This class is the base class for different mcrx shooting
    stages. It's derived from by the concrete stage classes, using the
    Barton-Nackman trick to make the types known to the generic
    routines. The definitions of the member functions are in the file
    mcrx-stage-impl.h, since they are only needed in the concrete
    stage classes that call them. */
template <template<typename> class grid_type, 
	  typename stage_type> class mcrx::mcrx_stage {
public:
  typedef typename stage_types<grid_type, stage_type>::T_grid T_grid;
  typedef typename stage_types<grid_type, stage_type>::T_emission T_emission;
  typedef typename stage_types<grid_type, stage_type>::T_emergence T_emergence;
  typedef typename stage_types<grid_type, stage_type>::T_dust_model T_dust_model;
  typedef typename stage_types<grid_type, stage_type>::T_dust_model::T_biaser T_biaser;
  typedef typename stage_types<grid_type, stage_type>::T_shooter T_shooter;
  typedef typename T_grid::T_data T_data;

private:
  stage_type& stage() {return static_cast<stage_type&>(*this);};

protected:

  Mcrx& m_; ///< Reference to parent mcrx object.
  Preferences& p_; ///< Reference to Preferences object in parent mcrx object.

  boost::shared_ptr<CCfits::FITS> input_file_, output_file_;
  boost::shared_ptr<T_grid> g_;
  boost::shared_ptr<T_emission> emi_;
  boost::shared_ptr<T_emergence> eme_;
  boost::shared_ptr<T_dust_model> model_;

  long n_rays_, n_rays_desired_;

  void read_units(T_unit_map&);
  /// Helper function that sets up a nonscatter shooting.
  bool shoot_nonscatter ();
  /** Helper function that sets up a scatter shooting. Add_intensity
      determines whether intensities should be tracked in the grid.
      This function is fairly trivial, since all the multithreading
      functionality has been wrapped into the general shoot
      function.  */
  bool shoot_scatter (bool add_intensity);
  
  long get_rays_completed ();
  void update_rays_completed (long);


public:
  mcrx_stage(Mcrx& m) : m_(m), p_(m.p) {};

  void run_stage();
};


// *** Aux grid stage ***

/** This is a traits class that defines the types needed for the
    aux_grid_stage. */
template <template<typename> class grid_type> 
class mcrx::stage_types<grid_type, mcrx::aux_grid_stage<grid_type> > {
public:
  typedef typename mcrx::aux_grid<grid_type, cumulative_sampling, local_random> T_grid;
  typedef typename mcrx::aux_grid<grid_type, cumulative_sampling, local_random> T_emission;
  typedef aux_grid_emergence T_emergence;
  typedef typename mcrx::generic_dust_model<typename T_emission::T_lambda> T_dust_model;
  typedef typename T_dust_model::T_biaser T_biaser;
  typedef typename mcrx::nonscatter_shooter<T_biaser> T_shooter;
};

/** This class performs the aux_grid_stage. It contains the
    definitions of the functions that are specific to this stage. */
template <template<typename> class grid_type>
class mcrx::aux_grid_stage : private mcrx_stage<grid_type, 
						aux_grid_stage<grid_type> > {
private:
  typedef mcrx_stage<grid_type, aux_grid_stage<grid_type> > T_base;
  typedef typename T_base::T_grid T_grid;
  typedef typename T_base::T_emission T_emission;
  typedef typename T_base::T_emergence T_emergence;
  typedef typename T_base::T_dust_model T_dust_model;
  typedef typename T_base::T_biaser T_biaser;
  typedef typename T_base::T_shooter T_shooter;
  friend class mcrx_stage<grid_type, aux_grid_stage<grid_type> >;

  std::string stage_ID () const {return "AUXG";};
  std::string stage_name () const {return "aux grid";};
  long get_rays_desired () const {return this->p_.getValue("nrays_aux", int());};
  long get_rays_completed ();
  void update_rays_completed (long);
  bool check_stage_state () const;
  void set_stage_state (bool);
  void setup_objects();
  void load_file ();
  void load_dump (binifstream&);
  void save_file ();
  void save_dump (binofstream&) const;
  bool shoot ();
  T_shooter shooter() {return T_shooter();};

public:
  aux_grid_stage(Mcrx& m) : T_base(m) {};

  void operator()();

};


// *** Aux particle stage ***

/** This is a traits class that defines the types needed for the
    aux_particle_stage. */
template <> class mcrx::stage_types<mcrx::dummy_grid_template, mcrx::aux_particle_stage> {
public:
  typedef dummy_grid T_grid;
  typedef aux_particles_emission<cumulative_sampling, local_random> T_emission;
  typedef aux_particle_emergence T_emergence;
  typedef generic_dust_model<T_emission::T_lambda> T_dust_model;
  typedef T_dust_model::T_biaser T_biaser;
  typedef nonscatter_shooter<T_biaser> T_shooter;
};

/** This class performs the aux_particle_stage. It contains the
    definitions of the functions that are specific to this stage. */
class mcrx::aux_particle_stage : private mcrx_stage<dummy_grid_template, aux_particle_stage> {
private:
  friend class mcrx_stage<dummy_grid_template, aux_particle_stage>;

  std::string stage_ID () const {return "AUXP";};
  std::string stage_name () const {return "aux particle";};
  long get_rays_desired () const {return p_.getValue("nrays_aux", int());};
  long get_rays_completed ();
  void update_rays_completed (long);
  bool check_stage_state () const;
  void set_stage_state (bool);
  void setup_objects();
  void load_file ();
  void load_dump (binifstream&);
  void save_file ();
  void save_dump (binofstream&) const;
  bool shoot ();
  T_shooter shooter() {return T_shooter();};

public:
  aux_particle_stage(Mcrx& m) : mcrx_stage<dummy_grid_template, aux_particle_stage>(m) {};

  void operator()();

};

// *** Nonscatter stage ***

/** This is a traits class that defines the types needed for the
    nonscatter_stage. */
template <> class mcrx::stage_types<mcrx::dummy_grid_template, mcrx::nonscatter_stage> {
public:
  typedef dummy_grid T_grid;
  typedef full_sed_particles_emission<cumulative_sampling,
    mcrx_rng_policy> T_emission;
  typedef full_sed_emergence T_emergence;
  typedef scatterer<polychromatic_scatterer_policy, 
    mcrx_rng_policy> T_scatterer;
  // this is nonoptimal. Because the emission classes are templated on
  // chromatic policy and full_sed_emission is a polychromatic
  // emission, we can't use the generic dust model because it is a
  // generic chroamtic policy even though both have types array_1.
  typedef dust_model<T_scatterer, cumulative_sampling,
    mcrx_rng_policy> T_dust_model;
  typedef T_dust_model::T_biaser T_biaser;
  typedef nonscatter_shooter<T_biaser> T_shooter;
};

/** This class performs the nonscatter_stage. It contains the
    definitions of the functions that are specific to this stage. */
class mcrx::nonscatter_stage : private mcrx_stage<dummy_grid_template, nonscatter_stage> {
private:
  friend class mcrx_stage<dummy_grid_template, nonscatter_stage>;

  std::string stage_ID () const {return "NONS";};
  std::string stage_name () const {return "nonscatter";};
  long get_rays_desired () const {
    return p_.getValue("nrays_nonscatter", int());};
  long get_rays_completed ();
  void update_rays_completed (long);
  bool check_stage_state () const;
  void set_stage_state (bool);
  void setup_objects();
  void load_file ();
  void load_dump (binifstream&);
  void save_file ();
  void save_dump (binofstream&) const;
  bool shoot ();
  T_shooter shooter() {return T_shooter(reflambda_);};

  T_dust_model::T_lambda lambda_;
  int reflambda_;
public:
  nonscatter_stage(Mcrx& m) : mcrx_stage<dummy_grid_template, nonscatter_stage>(m) {};

  void operator()();

};

// *** Scatter stage ***

/** This is a traits class that defines the types needed for the
    scatter_stage. */
template <template<typename> class grid_type> 
class mcrx::stage_types<grid_type, mcrx::scatter_stage<grid_type> > {
public:
  typedef typename mcrx::full_sed_grid<grid_type> T_grid;
  typedef full_sed_particles_emission<cumulative_sampling,
    mcrx_rng_policy> T_emission;
  typedef full_sed_emergence T_emergence;
  typedef scatterer<polychromatic_scatterer_policy, 
    mcrx_rng_policy> T_scatterer;
  typedef dust_model<T_scatterer, cumulative_sampling,
    mcrx_rng_policy> T_dust_model;
  typedef T_dust_model::T_biaser T_biaser;
  typedef scatter_shooter<T_biaser> T_shooter;
};

/** This class performs the scatter_stage. It contains the
    definitions of the functions that are specific to this stage. */
template <template<typename> class grid_type>
class mcrx::scatter_stage : private mcrx_stage<grid_type, scatter_stage<grid_type> > {
private:
  typedef mcrx_stage<grid_type, scatter_stage<grid_type> > T_base;
  typedef typename T_base::T_grid T_grid;
  typedef typename T_base::T_emission T_emission;
  typedef typename T_base::T_emergence T_emergence;
  typedef typename stage_types<grid_type, scatter_stage<grid_type> >::T_scatterer T_scatterer;
  typedef typename T_base::T_dust_model T_dust_model;
  typedef typename T_base::T_biaser T_biaser;
  typedef typename T_base::T_shooter T_shooter;
  friend class mcrx_stage<grid_type, scatter_stage<grid_type> >;

  std::string stage_ID () const {return "SCAT";};
  std::string stage_name () const {return "scatter";};
  long get_rays_desired () const {
    return this->p_.getValue("nrays_scatter", int());};
  long get_rays_completed ();
  void update_rays_completed (long);
  bool check_stage_state () const;
  void set_stage_state (bool);
  void setup_objects();
  void load_file ();
  void load_dump (binifstream&);
  void save_file ();
  void save_dump (binofstream&) const;
  bool shoot ();
  T_shooter shooter() {return T_shooter(reflambda_);};

  int reflambda_;
  typename T_dust_model::T_lambda lambda_;
public:
  scatter_stage(Mcrx& m) : T_base(m) {};

  void operator()();
};


// *** Intensity stage ***

/** This is a traits class that defines the types needed for the
    intensity_stage. */
template <template<typename> class grid_type> 
class mcrx::stage_types<grid_type, mcrx::intensity_stage<grid_type> > {
public:
  typedef typename mcrx::full_sed_grid<grid_type> T_grid;
  typedef full_sed_particles_emission<cumulative_sampling,
    mcrx_rng_policy> T_emission;
  typedef full_sed_emergence T_emergence;
  typedef scatterer<polychromatic_scatterer_policy, 
    mcrx_rng_policy> T_scatterer;
  typedef dust_model<T_scatterer, cumulative_sampling,
    mcrx_rng_policy> T_dust_model;
  typedef T_dust_model::T_biaser T_biaser;
  typedef scatter_shooter<T_biaser> T_shooter;
};

/** This class performs the intensity_stage. It contains the
    definitions of the functions that are specific to this stage. The
    intensity stage is used to estimate the radiation intensities
    resulting from the primary radiation. It used to be part of the
    scatter stage but it's better to have it separate because it can
    be run without cameras and with lower wavelength resolution. 
*/
template <template<typename> class grid_type>
class mcrx::intensity_stage : private mcrx_stage<grid_type, intensity_stage<grid_type> > {
private:
  typedef mcrx_stage<grid_type, intensity_stage<grid_type> > T_base;
  typedef typename T_base::T_grid T_grid;
  typedef typename T_base::T_emission T_emission;
  typedef typename T_base::T_emergence T_emergence;
  typedef typename stage_types<grid_type, intensity_stage<grid_type> >::T_scatterer T_scatterer;
  typedef typename T_base::T_dust_model T_dust_model;
  typedef typename T_base::T_biaser T_biaser;
  typedef typename T_base::T_shooter T_shooter;
  friend class mcrx_stage<grid_type, intensity_stage<grid_type> >;

  std::string stage_ID () const {return "INT";};
  std::string stage_name () const {return "intensity";};
  long get_rays_desired () const {
    return this->p_.getValue("nrays_intensity", int());};
  long get_rays_completed ();
  void update_rays_completed (long);
  bool check_stage_state () const;
  void set_stage_state (bool);
  void setup_objects();
  void load_file ();
  void load_dump (binifstream&);
  void save_file ();
  void save_dump (binofstream&) const;
  bool shoot ();
  T_shooter shooter() {return T_shooter(reflambda_);};

  int reflambda_;
  typename T_dust_model::T_lambda lambda_;
public:
  intensity_stage(Mcrx& m) : T_base(m) {};

  void operator()();

};


// *** IR stage ***

/** This is a traits class that defines the types needed for the
    ir_stage. */
template <template<typename> class grid_type>
class mcrx::stage_types<grid_type, mcrx::ir_stage<grid_type> > {
public:
  typedef typename mcrx::ir_grid<grid_type, cumulative_sampling, local_random> T_grid;
  typedef T_grid T_emission;
  typedef full_sed_emergence T_emergence;
  typedef scatterer<polychromatic_scatterer_policy, 
    mcrx_rng_policy> T_scatterer;
  // this is nonoptimal. Because the emission classes are templated on
  // chromatic policy and full_sed_emission is a polychromatic
  // emission, we can't use the generic dust model because it is a
  // generic chroamtic policy even though both have types array_1.
  typedef dust_model<T_scatterer, cumulative_sampling,
    mcrx_rng_policy> T_dust_model;
  typedef T_dust_model::T_biaser T_biaser;
  typedef nonscatter_shooter<T_biaser> T_shooter;
};

/** This class performs the ir_stage. It contains the
    definitions of the functions that are specific to this stage. */
template <template<typename> class grid_type>
class mcrx::ir_stage : private mcrx_stage<grid_type, ir_stage<grid_type> > {
private:
  typedef mcrx_stage<grid_type, ir_stage<grid_type> > T_base;
  typedef typename T_base::T_grid T_grid;
  typedef typename T_base::T_emission T_emission;
  typedef typename T_base::T_emergence T_emergence;
  typedef typename stage_types<grid_type, ir_stage<grid_type> >::T_scatterer T_scatterer;
  typedef typename T_base::T_dust_model T_dust_model;
  typedef typename T_base::T_biaser T_biaser;
  typedef typename T_base::T_shooter T_shooter;
  friend class mcrx_stage<grid_type, ir_stage<grid_type> >;

  std::string stage_ID () const {return "IR";};
  std::string stage_name () const {return "IR";};
  long get_rays_desired () const {
    return this->p_.getValue("nrays_ir", int());};
  long get_rays_completed ();
  void update_rays_completed (long);
  bool check_stage_state () const;
  void set_stage_state (bool);
  void setup_objects();
  void load_file ();
  void load_dump (binifstream&);
  void save_file ();
  void save_dump (binofstream&) const;
  bool shoot ();
  T_shooter shooter() {return T_shooter();};

  void basic_setup();
  void calculate_dust_SED();

  typename T_dust_model::T_lambda ilambda_, elambda_;
public:
  ir_stage(Mcrx& m) : T_base(m) {};

  void operator()();

};
  
// *** IR intensity stage ***

/** This is a traits class that defines the types needed for the
    thick_ir_stage.  Essentially, it has the same emission as the ir
    stage and the same grid as the scatter stage. */
template <template<typename> class grid_type>
class mcrx::stage_types<grid_type, mcrx::ir_intensity_stage<grid_type> > {
public:
  typedef typename mcrx::full_sed_grid<grid_type> T_grid;
  typedef typename mcrx::ir_grid<grid_type, cumulative_sampling, local_random> T_emission;
  typedef full_sed_emergence T_emergence;
  typedef scatterer<polychromatic_scatterer_policy, 
    mcrx_rng_policy> T_scatterer;
  typedef dust_model<T_scatterer, cumulative_sampling,
    mcrx_rng_policy> T_dust_model;
  typedef T_dust_model::T_biaser T_biaser;
  typedef scatter_shooter<T_biaser> T_shooter;
};

/** This class performs the ir_intensity_stage, which calculates the
    radiation intensity in the cells from dust emission. It contains
    the definitions of the functions that are specific to this
    stage. */
template <template<typename> class grid_type>
class mcrx::ir_intensity_stage : private mcrx_stage<grid_type, ir_intensity_stage<grid_type> > {
private:
  typedef mcrx_stage<grid_type, ir_intensity_stage<grid_type> > T_base;
  typedef typename T_base::T_grid T_grid;
  typedef typename T_base::T_emission T_emission;
  typedef typename T_base::T_emergence T_emergence;
  typedef typename stage_types<grid_type, ir_intensity_stage<grid_type> >::T_scatterer T_scatterer;
  typedef typename T_base::T_dust_model T_dust_model;
  typedef typename T_base::T_biaser T_biaser;
  typedef typename T_base::T_shooter T_shooter;
  friend class mcrx_stage<grid_type, ir_intensity_stage<grid_type> >;

  std::string stage_ID () const {return "IRIN";};
  std::string stage_name () const {return "IR Intensity";};
  long get_rays_desired () const {
    return this->p_.getValue("nrays_intensity", int());};
  long get_rays_completed ();
  void update_rays_completed (long);
  bool check_stage_state () const;
  void set_stage_state (bool);
  void setup_objects();
  void load_file ();
  void load_dump (binifstream&);
  void save_file ();
  void save_dump (binofstream&) const;
  bool shoot ();
  T_shooter shooter() {return T_shooter(reflambda_);};

  void basic_setup();
  template<typename T>
  void calculate_dust_SED(const blitz::ETBase<T>&);

  int reflambda_;
  typename T_dust_model::T_lambda lambda_;

  /// This vector keeps the same pointers as model_, but these are
  /// grain_model* and are used for the SED calculation.
  std::vector<grain_model<polychromatic_scatterer_policy, 
			  mcrx_rng_policy>*> grain_models_;
  array_2 star_intensity_, dust_intensity_;

public:
  ir_intensity_stage(Mcrx& m) : T_base(m) {};

  void operator()();

};


// *** optically thick IR stage ***

/** This is a traits class that defines the types needed for the
    thick_ir_stage.  Essentially, it has the same emission as the ir
    stage and the same grid as the scatter stage. */
template <template<typename> class grid_type>
class mcrx::stage_types<grid_type, mcrx::thick_ir_stage<grid_type> > {
public:
  typedef typename mcrx::full_sed_grid<grid_type> T_grid;
  typedef typename mcrx::ir_grid<grid_type, cumulative_sampling, local_random> T_emission;
  typedef full_sed_emergence T_emergence;
  typedef scatterer<polychromatic_scatterer_policy, 
    mcrx_rng_policy> T_scatterer;
  typedef dust_model<T_scatterer, cumulative_sampling,
    mcrx_rng_policy> T_dust_model;
  typedef T_dust_model::T_biaser T_biaser;
  typedef scatter_shooter<T_biaser> T_shooter;
};

/** This class performs the ir_stage. It contains the
    definitions of the functions that are specific to this stage. */
template <template<typename> class grid_type>
class mcrx::thick_ir_stage : private mcrx_stage<grid_type, thick_ir_stage<grid_type> > {
private:
  typedef mcrx_stage<grid_type, thick_ir_stage<grid_type> > T_base;
  typedef typename T_base::T_grid T_grid;
  typedef typename T_base::T_emission T_emission;
  typedef typename T_base::T_emergence T_emergence;
  typedef typename stage_types<grid_type, thick_ir_stage<grid_type> >::T_scatterer T_scatterer;
  typedef typename T_base::T_dust_model T_dust_model;
  typedef typename T_base::T_biaser T_biaser;
  typedef typename T_base::T_shooter T_shooter;
  friend class mcrx_stage<grid_type, thick_ir_stage<grid_type> >;

  std::string stage_ID () const {return "IR";};
  std::string stage_name () const {return "Optically thick IR";};
  long get_rays_desired () const {
    return this->p_.getValue("nrays_ir", int());};
  long get_rays_completed ();
  void update_rays_completed (long);
  bool check_stage_state () const;
  void set_stage_state (bool);
  void setup_objects();
  void load_file ();
  void load_dump (binifstream&);
  void save_file ();
  void save_dump (binofstream&) const;
  bool shoot ();
  T_shooter shooter() {return T_shooter(reflambda_);};

  void basic_setup();
  template<typename T>
  void calculate_dust_SED(const blitz::ETBase<T>&);

  int reflambda_;
  typename T_dust_model::T_lambda ilambda_, elambda_;

  /// This vector keeps the same pointers as model_, but these are
  /// grain_model* and are used for the SED calculation.
  std::vector<grain_model<polychromatic_scatterer_policy, 
			  mcrx_rng_policy>*> grain_models_;
  array_2 intensity_;

public:
  thick_ir_stage(Mcrx& m) : T_base(m) {};

  void operator()();

};

  
#endif
